// ATASHD3D.cpp : Defines the exported functions for the DLL application.
//

#include "stdafx.h"
#include "ATASHD3D.h"

#include <tchar.h>
#include <process.h>
#include <d3d9.h>
#include <stdio.h>
#include "HookD3D.h"
#include "CaptorManager.h"

CCaptorManager g_mgr;

HANDLE g_htrdMain = NULL;
HWND g_hwndFocus = NULL;

BOOL CreateCaptor(RPB_CreateCaptor* pParameter);
BOOL DestroyCaptor(HCAPTOR hCaptor);

#ifdef _DEBUG
void DebugInfoTempA(LPCSTR lpFormat, ...)
{
	CHAR tchBuf[1024] = {0};
	sprintf_s(tchBuf, 1024, "D3D Hook[%ld, %ld]", GetCurrentProcessId(), GetCurrentThreadId());

	va_list vl;
	CHAR tchMsg[1024] = {0};
	va_start(vl, lpFormat);
	vsprintf_s(tchMsg, 1024, lpFormat, vl);
	va_end(vl);

	strcat_s(tchBuf, 1024, tchMsg);
	strcat_s(tchBuf, 1024, "\n");
	OutputDebugStringA(tchBuf);
}

void DebugInfoTempW(LPCWSTR lpFormat, ...)
{
	WCHAR tchBuf[1024] = {0};
	swprintf_s(tchBuf, 1024, L"D3D Hook[%ld, %ld]", GetCurrentProcessId(), GetCurrentThreadId());

	va_list vl;
	WCHAR tchMsg[1024] = {0};
	va_start(vl, lpFormat);
	vswprintf_s(tchMsg, 1024, lpFormat, vl);
	va_end(vl);

	wcscat_s(tchBuf, 1024, tchMsg);
	wcscat_s(tchBuf, 1024, L"\n");
	OutputDebugStringW(tchBuf);
}
#endif

void BitBltBox32(LPBYTE pDestBmpData, int nDestWidth, int nDestHeight, 
				int nDestX, int nDestY , 
				LPBYTE pSrcBmpData, int nSrcWidth, int nSrcHeight, 
				int nSrcX, int nSrcY, int nCpyWidth, int nCpyHeight)
{	
	//D3D texture date is not upside down.
	int	nMemLineSize = (nDestWidth<<2);
	int	nBmpLineSize = (nSrcWidth<<2);

//	LPBYTE pDstData = pDestBmpData + (nDestHeight-(nDestY+nCpyHeight))*nMemLineSize + 
//		(nDestX<<2);
//	LPBYTE pSrcData = pSrcBmpData + (nSrcHeight - (nSrcY+nCpyHeight))*nBmpLineSize 
//		+ (nSrcX<<2);

	LPBYTE pDstData = pDestBmpData + nDestY*nMemLineSize + (nDestX<<2);
	LPBYTE pSrcData = pSrcBmpData + nSrcY*nBmpLineSize + (nSrcX<<2);

//	DebugInfoTempW(_T("BitBltBox --> Bmp(%d,%d)(%x, %x){%d, %d, %d, %d} ==> Cpy(%d,%d)(%x, %x){%d, %d, %d, %d}"),
//			nSrcWidth, nSrcHeight, pSrcBmpData, pSrcData, nSrcX, nSrcY, nCpyWidth, nCpyHeight, 
//			nDestWidth, nDestHeight, pDestBmpData, pDstData, nDestX, nDestY, nCpyWidth, nCpyHeight);

	int j = 0;
	int nRectCopyLen = nCpyWidth<<2;
	for(j=0; j<nCpyHeight; j++)
	{
		memcpy(pDstData, pSrcData, nRectCopyLen);
		pDstData += nMemLineSize;
		pSrcData += nBmpLineSize;
	}
}

BOOL CreateBMPFile(LPTSTR pszFile, PBITMAPINFO pbi, LPBYTE lpBits)
{ 
	HANDLE hf;                 // file handle 
	BITMAPFILEHEADER hdr;       // bitmap file-header 
	PBITMAPINFOHEADER pbih;     // bitmap info-header 
	DWORD dwTotal;              // total count of bytes 
	DWORD cb;                   // incremental count of bytes 
	BYTE *hp;                   // byte pointer 
	DWORD dwTmp; 

	pbih = (PBITMAPINFOHEADER) pbi; 

	// Create the .BMP file. 
	hf = CreateFile(pszFile, 
		GENERIC_READ | GENERIC_WRITE, 
		(DWORD) 0, 
		NULL, 
		CREATE_ALWAYS, 
		FILE_ATTRIBUTE_NORMAL, 
		(HANDLE) NULL); 
	if (hf == INVALID_HANDLE_VALUE) 
		return FALSE;

	hdr.bfType = 0x4d42;        // 0x42 = "B" 0x4d = "M" 
	// Compute the size of the entire file. 
	hdr.bfSize = (DWORD) (sizeof(BITMAPFILEHEADER) + 
		pbih->biSize + pbih->biClrUsed 
		* sizeof(RGBQUAD) + pbih->biSizeImage); 
	hdr.bfReserved1 = 0; 
	hdr.bfReserved2 = 0; 

	// Compute the offset to the array of color indices. 
	hdr.bfOffBits = (DWORD) sizeof(BITMAPFILEHEADER) + 
		pbih->biSize + pbih->biClrUsed 
		* sizeof (RGBQUAD); 

	// Copy the BITMAPFILEHEADER into the .BMP file. 
	if (!WriteFile(hf, (LPVOID) &hdr, sizeof(BITMAPFILEHEADER), 
		(LPDWORD) &dwTmp,  NULL)) 
		goto _end;

	// Copy the BITMAPINFOHEADER and RGBQUAD array into the file. 
	if (!WriteFile(hf, (LPVOID) pbih, sizeof(BITMAPINFOHEADER)+ pbih->biClrUsed * sizeof (RGBQUAD), (LPDWORD)&dwTmp, ( NULL))) 
		goto _end;

	// Copy the array of color indices into the .BMP file. 
	dwTotal = cb = pbih->biSizeImage; 
	hp = lpBits; 
	if (!WriteFile(hf, (LPSTR) hp, (int) cb, (LPDWORD) &dwTmp,NULL)) 
		goto _end;

	// Close the .BMP file. 
	if (!CloseHandle(hf)) 
		goto _end;

	return TRUE;
_end:
	if(hf!=INVALID_HANDLE_VALUE) CloseHandle(hf);
	return FALSE;
}

EXTERN_C IMAGE_DOS_HEADER __ImageBase;
#define HINST_THISCOMPONENT ((HINSTANCE)&__ImageBase)

LRESULT CALLBACK WindowProc(__in  HWND hwnd, __in  UINT uMsg, __in  WPARAM wParam, __in  LPARAM lParam)
{
	switch (uMsg)
	{
	case WM_CREATECAPTOR:
		CreateCaptor((RPB_CreateCaptor*)lParam);
		break;
	case WM_DESTROYCAPTOR:
		DestroyCaptor((HCAPTOR)lParam);
		break;
	case WM_DESTROY:
		PostQuitMessage(0);
		break;
	}
	return DefWindowProc(hwnd, uMsg, wParam, lParam);
}

void SaveTexturetoBmp(ID3D10Texture2D* pTexture, LPTSTR lpszFile)
{
	if(pTexture)
	{
		D3D10_TEXTURE2D_DESC textDesc;
		pTexture->GetDesc(&textDesc);

		DebugInfoTempA("Texture[%x, %d]:%d, %d, %d, %d, %d, %d, %d, %d, %d", 
			textDesc.CPUAccessFlags, textDesc.Usage, 
			textDesc.Height, textDesc.Width, textDesc.Format,
			textDesc.BindFlags,textDesc.MiscFlags, textDesc.MipLevels,
			textDesc.ArraySize, textDesc.SampleDesc.Count, textDesc.SampleDesc.Quality);

		ID3D10Device* pDevice = NULL;
		pTexture->GetDevice(&pDevice);

		UINT unViewPorts = 3;
		D3D10_VIEWPORT vp[3];
		pDevice->RSGetViewports(&unViewPorts, vp);
		DebugInfoTempA("ViewPort: %d vp: %d, %d, %d, %d", 
			unViewPorts, vp[0].TopLeftX, vp[0].TopLeftY, vp[0].Width, vp[0].Height);

		D3D10_RECT drect[3];
		UINT unScissorRect=3;
		pDevice->RSGetScissorRects(&unScissorRect, drect);
		DebugInfoTempA("ScissorRect: %d vp: %d, %d, %d, %d", 
			unScissorRect, drect[0].left, drect[0].top, 
			drect[0].right, drect[0].bottom);

		ID3D10Texture2D* pNewTexture = NULL;
		textDesc.CPUAccessFlags = D3D10_CPU_ACCESS_READ;
		textDesc.Usage = D3D10_USAGE_STAGING;
		textDesc.BindFlags = 0;
		HRESULT hr = pDevice->CreateTexture2D(&textDesc,NULL, &pNewTexture);

		pDevice->CopyResource(pNewTexture, pTexture);

		D3D10_MAPPED_TEXTURE2D mappedTex = {0};
		hr = pNewTexture->Map( D3D10CalcSubresource(0, 0, 1), D3D10_MAP_READ, 0, &mappedTex );

		BITMAPINFO bi;
		ZeroMemory(&bi, sizeof(bi));
		bi.bmiHeader.biSize = sizeof(BITMAPINFOHEADER);
		bi.bmiHeader.biBitCount = 32;
		bi.bmiHeader.biHeight = -(LONG)textDesc.Height;
		bi.bmiHeader.biWidth = textDesc.Width;
		bi.bmiHeader.biPlanes = 1;
		bi.bmiHeader.biCompression = BI_RGB;
		bi.bmiHeader.biSizeImage = mappedTex.RowPitch*textDesc.Height;

		TCHAR szFile[MAX_PATH];
		_stprintf_s(szFile, TEXT("d:\\temp\\%s"), lpszFile);
		CreateBMPFile(szFile, &bi, (LPBYTE)mappedTex.pData);
		pNewTexture->Unmap(D3D10CalcSubresource(0, 0, 1));


		DebugInfoTempA("Create new texture:%x, %x, %x, buf: %x, %d ", 
			pDevice, pNewTexture, hr, mappedTex.pData, mappedTex.RowPitch);

		if(pNewTexture) pNewTexture->Release();
		if(pDevice) pDevice->Release();
	}
}

DWORD WINAPI ThreadProc(__in  LPVOID lpParameter)
{
	MSG msg ={};
    BOOL bRet; 
    WNDCLASS wc; 
 
    // Register the window class for the main window. 
	TCHAR tchClass[128];
	_stprintf_s(tchClass, 128, TEXT("%s%d"), PREFIX_WINDOW_CLASS, GetCurrentProcessId());
	
    wc.style = 0; 
    wc.lpfnWndProc = (WNDPROC) WindowProc; 
    wc.cbClsExtra = 0; 
    wc.cbWndExtra = 0; 
    wc.hInstance = HINST_THISCOMPONENT; 
    wc.hIcon = NULL;
    wc.hCursor = LoadCursor((HINSTANCE) NULL, IDC_ARROW); 
    wc.hbrBackground = HBRUSH(COLOR_WINDOW + 1); 
    wc.lpszMenuName =  NULL; 
    wc.lpszClassName = tchClass; 

    if (!RegisterClass(&wc)) 
        return FALSE; 
 
    // Create the main window. 
 
    g_hwndFocus = CreateWindow(tchClass, NULL, 
        WS_POPUP, 0, 0, 
        1, 1, (HWND) NULL, 
        (HMENU) NULL, HINST_THISCOMPONENT, (LPVOID) NULL); 
 
    // If the main window cannot be created, terminate 
    // the application. 
 
    if (g_hwndFocus) 
	{
		g_mgr.Create10();

		// Start the message loop. 
 
		while( (bRet = GetMessage( &msg, NULL, 0, 0 )) != 0)
		{ 
			if (bRet == -1)
			{
				// handle the error and possibly exit
			}
			else
			{
				TranslateMessage(&msg); 
				DispatchMessage(&msg); 
			}
		} 
 
		// Return the exit code to the system. 
 
		g_mgr.Destroy();
	}

	UnregisterClass(tchClass, HINST_THISCOMPONENT);
    return msg.wParam; 
}

BOOL Initialize(PVOID pParameter)
{
	//g_mgr.Create();
	g_htrdMain = CreateThread(NULL, 0, ThreadProc, NULL, 0, NULL);
	return TRUE;
}

ATASHD3D_API BOOL RTInitialize(PVOID pParameter)
{
	DebugInfoTempW(TEXT("Initialize"));

	_endthreadex(Initialize(pParameter));
	return TRUE;
}

BOOL Finalize(PVOID pParameter)
{
	//g_mgr.Destroy();
	PostMessage(g_hwndFocus, WM_CLOSE, 0, 0);
	WaitForSingleObject(g_htrdMain, INFINITE);
	CloseHandle(g_htrdMain);
	return TRUE;
}

ATASHD3D_API BOOL RTFinalize(PVOID pParameter)
{
	DebugInfoTempW(TEXT("Finalize"));

	_endthreadex(Finalize(pParameter));
	return TRUE;
}

BOOL CreateCaptor(RPB_CreateCaptor* pParameter)
{
	CCaptor* pCaptor = g_mgr.CreateCaptor(pParameter->szDeviceName);
	if (!pCaptor)
		return FALSE;

	RPB_CreateCaptor* p = pCaptor->GetCreateParameter();
	memcpy(pParameter, p, sizeof(RPB_CreateCaptor));

//	pParameter->hCaptor = (HCAPTOR)pCaptor;
//	pParameter->hMutex = pCaptor->GetMutex();
//	pParameter->hFileMapping = pCaptor->GetFileMapping();
//	pParameter->hWndNotified = g_hwndFocus;

	return TRUE;
}

ATASHD3D_API BOOL RTCreateCaptor(RPB_CreateCaptor* pParameter)
{
	DebugInfoTempA("CreateCaptor: %s", pParameter->szDeviceName);
	_endthreadex(CreateCaptor(pParameter));
	return 2;
}

BOOL DestroyCaptor(HCAPTOR hCaptor)
{
	g_mgr.DestroyCaptor((CCaptor*)hCaptor);
	return TRUE;
}

ATASHD3D_API BOOL RTDestroyCaptor(HCAPTOR hCaptor)
{
	DebugInfoTempW(TEXT("DestroyCaptor(%08X)"), hCaptor);
	_endthreadex(DestroyCaptor(hCaptor));
	return TRUE;
}
